package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/fatih/color"
	"github.com/olekukonko/tablewriter"
)

// ServerConfig 定义服务器配置结构体
type ServerConfig struct {
	Name             string `json:"name"`
	URL              string `json:"url"`
	Username         string `json:"username"`
	Password         string `json:"password"`
	DomainBruteType  string `json:"domain_brute_type"`
	PortScanType     string `json:"port_scan_type"`
	DomainBrute      bool   `json:"domain_brute"`
	AltDNS           bool   `json:"alt_dns"`
	DNSQueryPlugin   bool   `json:"dns_query_plugin"`
	ARLSearch        bool   `json:"arl_search"`
	PortScan         bool   `json:"port_scan"`
	ServiceDetection bool   `json:"service_detection"`
	OSDetection      bool   `json:"os_detection"`
	SSLCert          bool   `json:"ssl_cert"`
	SkipScanCDNIP    bool   `json:"skip_scan_cdn_ip"`
	SiteIdentify     bool   `json:"site_identify"`
	SearchEngines    bool   `json:"search_engines"`
	SiteSpider       bool   `json:"site_spider"`
	SiteCapture      bool   `json:"site_capture"`
	FileLeak         bool   `json:"file_leak"`
	FindVhost        bool   `json:"findvhost"`
	NucleiScan       bool   `json:"nuclei_scan"`
	WebInfoHunter    bool   `json:"web_info_hunter"`
}

// Config 定义配置文件结构体
type Config struct {
	Servers []ServerConfig `json:"servers"`
}

// TaskListResponse 定义任务列表响应结构体
type TaskListResponse struct {
	Page  int `json:"page"`
	Size  int `json:"size"`
	Total int `json:"total"`
	Items []struct {
		ID        string `json:"_id"`
		Name      string `json:"name"`
		Target    string `json:"target"`
		Status    string `json:"status"`
		StartTime string `json:"start_time"`
		EndTime   string `json:"end_time"`
		Statistic struct {
			SiteCnt   int `json:"site_cnt"`
			DomainCnt int `json:"domain_cnt"`
		} `json:"statistic"`
	} `json:"items"`
}

// LoginResponse 定义登录响应结构体
type LoginResponse struct {
	Code int `json:"code"`
	Data struct {
		Token string `json:"token"`
	} `json:"data"`
}

// Task 定义任务结构体
type Task struct {
	Name             string `json:"name"`
	Target           string `json:"target"`
	DomainBruteType  string `json:"domain_brute_type"`
	PortScanType     string `json:"port_scan_type"`
	DomainBrute      bool   `json:"domain_brute"`
	AltDNS           bool   `json:"alt_dns"`
	DNSQueryPlugin   bool   `json:"dns_query_plugin"`
	ARLSearch        bool   `json:"arl_search"`
	PortScan         bool   `json:"port_scan"`
	ServiceDetection bool   `json:"service_detection"`
	OSDetection      bool   `json:"os_detection"`
	SSLCert          bool   `json:"ssl_cert"`
	SkipScanCDNIP    bool   `json:"skip_scan_cdn_ip"`
	SiteIdentify     bool   `json:"site_identify"`
	SearchEngines    bool   `json:"search_engines"`
	SiteSpider       bool   `json:"site_spider"`
	SiteCapture      bool   `json:"site_capture"`
	FileLeak         bool   `json:"file_leak"`
	FindVhost        bool   `json:"findvhost"`
	NucleiScan       bool   `json:"nuclei_scan"`
	WebInfoHunter    bool   `json:"web_info_hunter"`
}

const (
	configFile = "config.json"
	logFile    = "wenarl.log"
)

var logger *log.Logger
var httpClient *http.Client

// 初始化日志记录器和HTTP客户端
func init() {
	// Initialize logger
	file, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		fmt.Printf("无法打开日志文件 %s: %v\n", logFile, err)
		os.Exit(1)
	}
	logger = log.New(file, "", log.Ldate|log.Ltime|log.Lshortfile)

	// Initialize HTTP client
	httpClient = &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		},
		Timeout: 30 * time.Second, // 设置超时时间
	}
}

func main() {
	// 显示启动字符画
	displayBanner()

	// 检查配置文件是否存在，如果不存在则创建默认配置文件并提示用户修改
	config, err := loadOrCreateConfig(configFile)
	if err != nil {
		color.Red("[-] 😒 读取配置文件失败: %v", err)
		logger.Printf("读取配置文件失败: %v", err)
		return
	}

	// 如果配置文件是新创建的，提醒用户修改配置文件并退出
	if config == nil {
		color.Yellow("请先修改配置文件 %s，然后重新运行程序。", configFile)
		logger.Printf("配置文件不存在，已创建默认配置文件 %s。", configFile)
		return
	}

	// 验证配置
	if err := validateConfig(config); err != nil {
		color.Red("[-] 配置验证失败: %v", err)
		logger.Printf("配置验证失败: %v", err)
		return
	}

	// 解析命令行参数
	taskCmd := flag.Bool("task", false, "输出当前ARL的任务状态")
	serversCmd := flag.Bool("servers", false, "启用多服务器模式进行任务分配")
	taskName := flag.String("n", "", "任务名称 (必填)")
	target := flag.String("t", "", "目标IP或域名")
	filePath := flag.String("f", "", "从文件读取扫描目标")
	outCmd := flag.String("out", "", "导出指定任务ID的报告（多服务器模式下可用逗号分隔多个ID）")
	outAllCmd := flag.Bool("outall", false, "导出所有指定任务名称的报告（多服务器模式下）")
	servers := flag.String("s", "", "指定服务器别名（多服务器模式下可用逗号分隔多个别名）")
	numTasks := flag.Int("num", 30, "指定查看的任务数量（与 -task 一起使用，默认30，最大20）")
	flag.Parse()

	// 限制 -num 参数最大值为20
	if *numTasks > 20 {
		*numTasks = 20
	}

	// 登录所有服务器并获取Token
	var tokens []string
	var serverNames []string
	selectedServers := config.Servers

	if *serversCmd || *servers != "" {
		// 多服务器模式
		var selectedServerConfigs []ServerConfig
		if *serversCmd {
			selectedServerConfigs = config.Servers
		} else if *servers != "" {
			aliases := strings.Split(*servers, ",")
			for _, alias := range aliases {
				alias = strings.TrimSpace(alias)
				found := false
				for _, server := range config.Servers {
					if server.Name == alias {
						selectedServerConfigs = append(selectedServerConfigs, server)
						found = true
						break
					}
				}
				if !found {
					color.Red("[-] 未找到别名为 %s 的服务器配置", alias)
					logger.Printf("未找到别名为 %s 的服务器配置", alias)
				}
			}
		}

		if len(selectedServerConfigs) == 0 {
			color.Red("[-] 未选择任何有效的服务器")
			logger.Println("未选择任何有效的服务器")
			return
		}

		selectedServers = selectedServerConfigs

		// 预先分配 tokens 和 serverNames 列表，确保索引对应
		tokens = make([]string, len(selectedServers))
		serverNames = make([]string, len(selectedServers))

		var wg sync.WaitGroup
		var mu sync.Mutex
		var errorsList []error

		for i, server := range selectedServers {
			wg.Add(1)
			go func(i int, server ServerConfig) {
				defer wg.Done()
				token, err := login(server.URL, server.Username, server.Password)
				if err != nil {
					mu.Lock()
					errorsList = append(errorsList, fmt.Errorf("服务器 %s 登录失败: %v", server.Name, err))
					mu.Unlock()
					color.Red("[-] %s 登录失败: %v", server.Name, err)
					logger.Printf("服务器 %s 登录失败: %v", server.Name, err)
					return
				}
				tokens[i] = token
				serverNames[i] = server.Name
				color.Green("[+] 🕵️ %s 登录成功！", server.Name)
				logger.Printf("服务器 %s 登录成功", server.Name)
			}(i, server)
		}

		wg.Wait()

		// 过滤成功登录的服务器和Token
		activeServers := []ServerConfig{}
		activeTokens := []string{}
		activeServerNames := []string{}

		for i, token := range tokens {
			if token != "" {
				activeServers = append(activeServers, selectedServers[i])
				activeTokens = append(activeTokens, token)
				activeServerNames = append(activeServerNames, serverNames[i])
			}
		}

		if len(activeTokens) == 0 {
			color.Red("[-] 所有选定的服务器登录失败，程序退出。")
			logger.Println("所有选定的服务器登录失败")
			return
		}

		if len(activeTokens) < len(selectedServers) {
			color.Yellow("部分服务器登录失败，继续处理其他服务器。")
			for _, e := range errorsList {
				color.Yellow("  - %v", e)
				logger.Printf("部分服务器登录失败: %v", e)
			}
		}

		// 更新 selectedServers 和 tokens 为登录成功的服务器
		selectedServers = activeServers
		tokens = activeTokens
		serverNames = activeServerNames
	} else {
		// 单服务器模式，使用第一个服务器
		firstServer := config.Servers[0]
		token, err := login(firstServer.URL, firstServer.Username, firstServer.Password)
		if err != nil {
			color.Red("[-] %s 登录失败: %v", firstServer.Name, err)
			logger.Printf("服务器 %s 登录失败: %v", firstServer.Name, err)
			return
		}
		tokens = append(tokens, token)
		serverNames = append(serverNames, firstServer.Name)
		color.Green("[+] 🕵️ %s 登录成功！", firstServer.Name)
		logger.Printf("服务器 %s 登录成功", firstServer.Name)
	}

	// 如果指定了 -task 参数，输出当前任务状态
	if *taskCmd {
		if len(selectedServers) > 1 && (*serversCmd || *servers != "") {
			// 多服务器模式
			var wg sync.WaitGroup
			for i, server := range selectedServers {
				if i >= len(tokens) {
					continue
				}
				wg.Add(1)
				go func(server ServerConfig, token string) {
					defer wg.Done()
					color.Cyan("\n🔹 %s 当前任务状态:", server.Name)
					getTaskList(server.URL, token, *numTasks)
				}(server, tokens[i])
			}
			wg.Wait()
		} else {
			// 单服务器模式
			firstServer := selectedServers[0]
			color.Cyan("\n🔹 %s 当前任务状态:", firstServer.Name)
			getTaskList(firstServer.URL, tokens[0], *numTasks)
		}
		return
	}

	// 如果指定了 -out 参数，导出任务报告
	if *outCmd != "" {
		if len(selectedServers) > 1 && (*serversCmd || *servers != "") {
			// 多服务器模式，支持多个任务ID
			taskIDs := strings.Split(*outCmd, ",")
			color.Cyan("开始导出任务报告，任务ID: %s", *outCmd)
			logger.Printf("开始导出任务报告，任务ID: %s", *outCmd)

			var wg sync.WaitGroup
			var mu sync.Mutex
			var errorsList []error

			for i, taskID := range taskIDs {
				taskID = strings.TrimSpace(taskID)
				if i >= len(tokens) {
					break
				}
				server := selectedServers[i%len(selectedServers)]
				token := tokens[i%len(tokens)]
				wg.Add(1)
				go func(server ServerConfig, token, taskID string) {
					defer wg.Done()
					err := exportTaskReport(server.URL, token, taskID)
					if err != nil {
						mu.Lock()
						errorsList = append(errorsList, fmt.Errorf("服务器 %s 任务ID %s 导出失败: %v", server.Name, taskID, err))
						mu.Unlock()
						color.Red("[-] %s 任务ID %s 报告导出失败: %v", server.Name, taskID, err)
						logger.Printf("服务器 %s 任务ID %s 报告导出失败: %v", server.Name, taskID, err)
					} else {
						color.Green("[+] %s 任务ID %s 报告导出成功！", server.Name, taskID)
						logger.Printf("服务器 %s 任务ID %s 报告导出成功", server.Name, taskID)
					}
				}(server, token, taskID)
			}

			wg.Wait()

			if len(errorsList) > 0 {
				color.Red("部分报告导出失败:")
				for _, e := range errorsList {
					color.Red("  - %v", e)
					logger.Printf("报告导出失败: %v", e)
				}
			} else {
				color.Green("[+] 所有报告已成功导出")
				logger.Println("所有报告已成功导出")
			}
		} else {
			// 单服务器模式
			firstServer := selectedServers[0]
			color.Cyan("开始导出任务报告，任务ID: %s", *outCmd)
			logger.Printf("开始导出任务报告，任务ID: %s", *outCmd)
			err := exportTaskReport(firstServer.URL, tokens[0], *outCmd)
			if err != nil {
				color.Red("[-] 报告导出失败: %v", err)
				logger.Printf("服务器 %s 报告导出失败: %v", firstServer.Name, err)
				os.Exit(1)
			}
			color.Green("[+] 报告导出成功！")
			logger.Printf("服务器 %s 报告导出成功，任务ID: %s", firstServer.Name, *outCmd)
		}
		return
	}

	// 如果指定了 -outall 参数，导出所有指定任务名称的报告
	if *outAllCmd {
		if *taskName == "" {
			color.Red("[-] 使用 -outall 时，必须指定任务名称 -n <任务名称>")
			logger.Println("使用 -outall 时，未指定任务名称")
			os.Exit(1)
		}

		if len(selectedServers) > 1 && (*serversCmd || *servers != "") {
			// 多服务器模式，导出所有服务器上指定任务名称的报告
			color.Cyan("开始导出所有服务器上任务名称为 '%s' 的报告", *taskName)
			logger.Printf("开始导出所有服务器上任务名称为 '%s' 的报告", *taskName)

			var wg sync.WaitGroup
			var mu sync.Mutex
			var errorsList []error

			for i, server := range selectedServers {
				if i >= len(tokens) {
					continue
				}
				wg.Add(1)
				go func(server ServerConfig, token string) {
					defer wg.Done()
					err := exportTaskReportAll(server.URL, token, *taskName)
					if err != nil {
						mu.Lock()
						errorsList = append(errorsList, fmt.Errorf("服务器 %s 导出失败: %v", server.Name, err))
						mu.Unlock()
						color.Red("[-] %s 报告导出失败: %v", server.Name, err)
						logger.Printf("服务器 %s 报告导出失败: %v", server.Name, err)
					} else {
						color.Green("[+] %s 报告导出成功！", server.Name)
						logger.Printf("服务器 %s 报告导出成功，任务名称: %s", server.Name, *taskName)
					}
				}(server, tokens[i])
			}

			wg.Wait()

			if len(errorsList) > 0 {
				color.Red("部分报告导出失败:")
				for _, e := range errorsList {
					color.Red("  - %v", e)
					logger.Printf("报告导出失败: %v", e)
				}
			} else {
				color.Green("[+] 所有报告已成功导出")
				logger.Println("所有报告已成功导出")
			}
		} else {
			// 单服务器模式，导出指定任务名称的所有报告
			firstServer := selectedServers[0]
			color.Cyan("开始导出任务名称为 '%s' 的所有报告", *taskName)
			logger.Printf("开始导出任务名称为 '%s' 的所有报告", *taskName)
			err := exportTaskReportAll(firstServer.URL, tokens[0], *taskName)
			if err != nil {
				color.Red("[-] 报告导出失败: %v", err)
				logger.Printf("服务器 %s 报告导出失败: %v", firstServer.Name, err)
				os.Exit(1)
			}
			color.Green("[+] 报告导出成功！")
			logger.Printf("服务器 %s 报告导出成功，任务名称: %s", firstServer.Name, *taskName)
		}
		return
	}

	// 如果指定了 -f 参数，从文件读取目标
	if *filePath != "" {
		targets, err := readTargetsFromFile(*filePath)
		if err != nil {
			color.Red("读取文件失败: %v", err)
			logger.Printf("读取文件失败: %v", err)
			os.Exit(1) // 文件读取失败，退出程序
		}

		if *taskName == "" {
			color.Red("任务名称未指定，请使用 -n <任务名称>")
			logger.Println("任务名称未指定")
			os.Exit(1)
		}

		// 输出扫描配置信息一次
		if len(selectedServers) > 1 && (*serversCmd || *servers != "") {
			displayScanConfigs(selectedServers)
		} else {
			displayScanConfig(&selectedServers[0])
		}

		if len(selectedServers) > 1 && (*serversCmd || *servers != "") {
			distributeTasks(selectedServers, tokens, *taskName, targets)
		} else {
			// 单服务器模式，创建所有任务在第一个服务器
			var wg sync.WaitGroup
			var mu sync.Mutex
			var errorsList []error

			// 定义并发限制，例如每次最多10个并发任务
			const maxConcurrentTasks = 10
			semaphore := make(chan struct{}, maxConcurrentTasks)

			for _, tgt := range targets {
				wg.Add(1)
				go func(target string) {
					defer wg.Done()
					semaphore <- struct{}{} // 获取一个信号量
					defer func() { <-semaphore }() // 释放信号量

					err := addTask(selectedServers[0].URL, tokens[0], *taskName, target, &selectedServers[0])
					if err != nil {
						mu.Lock()
						errorsList = append(errorsList, fmt.Errorf("目标 %s: %v", target, err))
						mu.Unlock()
						color.Red("[-] 目标 %s 创建任务失败: %v", target, err)
						logger.Printf("目标 %s 创建任务失败: %v", target, err)
					} else {
						color.Green("[+] 任务已成功创建：%s 目标：%s", *taskName, target)
						logger.Printf("任务已成功创建：%s 目标：%s", *taskName, target)
					}
				}(tgt)
			}

			wg.Wait()

			if len(errorsList) > 0 {
				color.Red("部分任务创建失败:")
				for _, e := range errorsList {
					color.Red("  - %v", e)
					logger.Printf("任务创建失败: %v", e)
				}
			} else {
				color.Green("[+] 所有任务已成功创建")
				logger.Println("所有任务已成功创建")
			}
		}
		return
	}

	// 如果提供了任务名称和目标，则创建任务
	if *taskName != "" && *target != "" {
		if len(selectedServers) > 1 && (*serversCmd || *servers != "") {
			// 多服务器模式，分配任务到各服务器
			distributeTasks(selectedServers, tokens, *taskName, []string{*target})
		} else {
			// 单服务器模式
			err = addTask(selectedServers[0].URL, tokens[0], *taskName, *target, &selectedServers[0])
			if err != nil {
				color.Red("任务创建失败: %v", err)
				logger.Printf("任务创建失败: %v", err)
				os.Exit(1) // 任务创建失败，退出程序
			}
			color.Green("[+] 任务已成功创建：%s 目标：%s", *taskName, *target)
			logger.Printf("任务已成功创建：%s 目标：%s", *taskName, *target)

			// 输出扫描配置信息一次
			displayScanConfig(&selectedServers[0])
		}
		return
	}

	// 如果参数不足，显示使用方法
	showUsage()
	color.Red("请使用 -n <任务名称> 和 -t <目标> 或 -f <文件路径> 来创建任务，或使用 -task 查看任务状态")
	logger.Println("参数不足")
	os.Exit(1) // 参数不足，退出程序
}

// 显示启动字符画
func displayBanner() {
	styledBanner := `
 __      __                 ______         ___      
/\ \  __/\ \               /\  _  \       /\_ \     
\ \ \/\ \ \ \     __    ___\ \ \L\ \  _ __\//\ \    
 \ \ \ \ \ \ \  /'__` + "`" + `\/' _ ` + "`" + `\ \  __ \/\` + "`" + `'__\\ \ \   
  \ \ \_/ \_\ \/\  __//\ \/\ \ \ \/\ \ \ \/  \_\ \_ 
   \ ` + "`" + `\___x___/\ \____\ \_\ \_\ \_\ \_\ \_\  /\____\
    '\/__//__/  \/____/\/_/\/_/\/_/\/_/\/_/  \/____/
                                                           
                                        Author: 温茶 
`
	color.Magenta(styledBanner)
}

// 显示使用方法
func showUsage() {
	color.Cyan("Usage:")
	fmt.Println("  -servers                     启用多服务器模式进行任务分配")
	fmt.Println("  -s <别名1,别名2,...>        指定服务器别名（多服务器模式下可用逗号分隔多个别名）")
	fmt.Println("  -n <任务名称>                设置任务名称 (必填)")
	fmt.Println("  -t <目标>                    设置扫描目标IP或域名")
	fmt.Println("  -f <文件路径>                从文件读取扫描目标")
	fmt.Println("  -task                        查看当前任务状态")
	fmt.Println("  -out <任务ID1,任务ID2,...>    导出指定任务ID的报告（多服务器模式下可用逗号分隔多个ID）")
	fmt.Println("  -outall                      导出所有指定任务名称的报告（多服务器模式下）")
	fmt.Println("  -num <数量>                  指定查看的任务数量（与 -task 一起使用，默认30，最大20）")
	fmt.Println("\n示例:")
	fmt.Println("  wenarl -n mytask -t example.com 指定单个任务去扫描")
	fmt.Println("  wenarl -n mytask -f targets.txt 指定文件中的任务扫描")
	fmt.Println("  wenarl -task -num 20 查看单服务器 20条最近任务")
	fmt.Println("  wenarl -servers -n 多目标扫描 -f url.txt 使用多服务器模式从文件中读取目标资产 下发到多个服务器(不重复) ")
	fmt.Println("  wenarl -servers -s server1name,server2name -n 指定的任务名称 -out id1,id2 导出不同服务器的不同任务")
	fmt.Println("  wenarl -servers -s server1name,server2name -n 指定的任务名称 -outall 导出指定服务器内相同任务名的所有报告")
}

// 加载或创建配置文件
func loadOrCreateConfig(filename string) (*Config, error) {
	file, err := os.Open(filename)
	if os.IsNotExist(err) {
		// 如果文件不存在，创建默认配置文件
		defaultConfig := &Config{
			Servers: []ServerConfig{
				{
					Name:             "server1name",
					URL:              "https://server1.example.com/",
					Username:         "admin",
					Password:         "password1",
					DomainBruteType:  "big",
					PortScanType:     "all",
					DomainBrute:      true,
					AltDNS:           true,
					DNSQueryPlugin:   true,
					ARLSearch:        true,
					PortScan:         true,
					ServiceDetection: true,
					OSDetection:      true,
					SSLCert:          false,
					SkipScanCDNIP:    true,
					SiteIdentify:     true,
					SearchEngines:    true,
					SiteSpider:       true,
					SiteCapture:      true,
					FileLeak:         true,
					FindVhost:        true,
					NucleiScan:       false,
					WebInfoHunter:    false,
				},
				{
					Name:             "server2name",
					URL:              "https://server2.example.com/",
					Username:         "admin",
					Password:         "password2",
					DomainBruteType:  "big",
					PortScanType:     "all",
					DomainBrute:      true,
					AltDNS:           true,
					DNSQueryPlugin:   true,
					ARLSearch:        true,
					PortScan:         true,
					ServiceDetection: true,
					OSDetection:      true,
					SSLCert:          false,
					SkipScanCDNIP:    true,
					SiteIdentify:     true,
					SearchEngines:    true,
					SiteSpider:       true,
					SiteCapture:      true,
					FileLeak:         true,
					FindVhost:        true,
					NucleiScan:       false,
					WebInfoHunter:    false,
				},
			},
		}
		err = createDefaultConfig(filename, defaultConfig)
		if err != nil {
			return nil, err
		}
		color.Green("配置文件不存在，已创建默认配置文件 %s。", filename)
		logger.Printf("配置文件不存在，已创建默认配置文件 %s。", filename)
		return nil, nil // 提示用户后退出程序
	} else if err != nil {
		return nil, err
	}
	defer file.Close()

	// 加载现有配置文件
	var config Config
	err = json.NewDecoder(file).Decode(&config)
	if err != nil {
		return nil, err
	}
	return &config, nil
}

// 创建默认配置文件
func createDefaultConfig(filename string, config *Config) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	data, err := json.MarshalIndent(config, "", "  ")
	if err != nil {
		return err
	}

	_, err = file.Write(data)
	if err != nil {
		return err
	}
	return nil
}

// 登录函数，获取Token
func login(url, username, password string) (string, error) {
	loginData := map[string]string{
		"username": username,
		"password": password,
	}
	jsonData, err := json.Marshal(loginData)
	if err != nil {
		return "", fmt.Errorf("登录数据序列化失败: %v", err)
	}

	resp, err := httpClient.Post(url+"api/user/login", "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		return "", fmt.Errorf("发送登录请求失败,请检查服务器地址以及账号密码: %v")
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("登录响应错误: %s", string(body))
	}

	var loginResp LoginResponse
	err = json.Unmarshal(body, &loginResp)
	if err != nil {
		return "", fmt.Errorf("解析登录响应失败: %v", err)
	}

	if loginResp.Code != 200 {
		return "", fmt.Errorf("登录失败，响应代码: %d", loginResp.Code)
	}

	return loginResp.Data.Token, nil
}

// 从文件读取目标
func readTargetsFromFile(filePath string) ([]string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	content, err := ioutil.ReadAll(file)
	if err != nil {
		return nil, err
	}

	// 按行分割目标
	targets := strings.Split(string(content), "\n")
	// 去除空行
	var validTargets []string
	for _, target := range targets {
		trimmed := strings.TrimSpace(target)
		if trimmed != "" {
			validTargets = append(validTargets, trimmed)
		}
	}
	return validTargets, nil
}

// 创建任务
func addTask(url, token, taskName, target string, config *ServerConfig) error {
	task := Task{
		Name:             taskName,
		Target:           target,
		DomainBruteType:  config.DomainBruteType,
		PortScanType:     config.PortScanType,
		DomainBrute:      config.DomainBrute,
		AltDNS:           config.AltDNS,
		DNSQueryPlugin:   config.DNSQueryPlugin,
		ARLSearch:        config.ARLSearch,
		PortScan:         config.PortScan,
		ServiceDetection: config.ServiceDetection,
		OSDetection:      config.OSDetection,
		SSLCert:          config.SSLCert,
		SkipScanCDNIP:    config.SkipScanCDNIP,
		SiteIdentify:     config.SiteIdentify,
		SearchEngines:    config.SearchEngines,
		SiteSpider:       config.SiteSpider,
		SiteCapture:      config.SiteCapture,
		FileLeak:         config.FileLeak,
		FindVhost:        config.FindVhost,
		NucleiScan:       config.NucleiScan,
		WebInfoHunter:    config.WebInfoHunter,
	}

	jsonData, err := json.Marshal(task)
	if err != nil {
		return fmt.Errorf("任务序列化失败: %v", err)
	}

	req, err := http.NewRequest("POST", url+"api/task/", bytes.NewBuffer(jsonData))
	if err != nil {
		return fmt.Errorf("创建HTTP请求失败: %v", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Token", token)

	resp, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("发送请求失败: %v", err)
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		return fmt.Errorf("任务创建失败: %s", string(body))
	}

	// 记录日志
	logger.Printf("任务创建成功：%s 目标：%s", taskName, target)
	return nil
}

// 导出单个任务报告
func exportTaskReport(serverURL, token, taskID string) error {
	req, err := http.NewRequest("GET", serverURL+"api/export/"+taskID, nil)
	if err != nil {
		return fmt.Errorf("创建HTTP请求失败: %v", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Token", token)

	resp, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("发送请求失败: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		body, _ := ioutil.ReadAll(resp.Body)
		return fmt.Errorf("导出失败，响应: %s", string(body))
	}

	// 获取文件名并进行URL解码
	contentDisposition := resp.Header.Get("Content-Disposition")
	filename := "report.xlsx"
	if strings.Contains(contentDisposition, "filename=") {
		filename = strings.Split(contentDisposition, "filename=")[1]
		filename = strings.Trim(filename, "\"")
		decodedFilename, err := url.QueryUnescape(filename)
		if err == nil {
			filename = decodedFilename
		}
	}

	// 保存文件
	outFile, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("创建文件失败: %v", err)
	}
	defer outFile.Close()

	_, err = io.Copy(outFile, resp.Body)
	if err != nil {
		return fmt.Errorf("保存文件失败: %v", err)
	}

	color.Green("报告已成功导出为: %s", filename)
	logger.Printf("报告已成功导出为: %s", filename)
	return nil
}

// 导出所有指定任务名称的报告
func exportTaskReportAll(serverURL, token, taskName string) error {
	// 获取所有任务列表
	requestURL := fmt.Sprintf("%sapi/task/?size=1000&page=1", serverURL)
	req, err := http.NewRequest("GET", requestURL, nil)
	if err != nil {
		return fmt.Errorf("创建HTTP请求失败: %v", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Token", token)

	resp, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("发送请求失败: %v", err)
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		return fmt.Errorf("获取任务列表失败: %s", string(body))
	}

	// 解析任务列表
	var taskList TaskListResponse
	err = json.Unmarshal(body, &taskList)
	if err != nil {
		return fmt.Errorf("解析任务列表失败: %v", err)
	}

	// 找到所有匹配的任务ID
	var matchingTaskIDs []string
	for _, task := range taskList.Items {
		if task.Name == taskName {
			matchingTaskIDs = append(matchingTaskIDs, task.ID)
		}
	}

	if len(matchingTaskIDs) == 0 {
		return fmt.Errorf("没有找到任务名称为 '%s' 的任务", taskName)
	}

	// 导出每个任务ID的报告
	var wg sync.WaitGroup
	var mu sync.Mutex
	var errorsList []error

	for _, taskID := range matchingTaskIDs {
		wg.Add(1)
		go func(taskID string) {
			defer wg.Done()
			req, err := http.NewRequest("GET", serverURL+"api/export/"+taskID, nil)
			if err != nil {
				mu.Lock()
				errorsList = append(errorsList, fmt.Errorf("创建HTTP请求失败: %v", err))
				mu.Unlock()
				logger.Printf("创建HTTP请求失败: %v", err)
				return
			}

			req.Header.Set("Content-Type", "application/json")
			req.Header.Set("Token", token)

			resp, err := httpClient.Do(req)
			if err != nil {
				mu.Lock()
				errorsList = append(errorsList, fmt.Errorf("发送请求失败: %v", err))
				mu.Unlock()
				logger.Printf("发送请求失败: %v", err)
				return
			}
			defer resp.Body.Close()

			if resp.StatusCode != 200 {
				body, _ := ioutil.ReadAll(resp.Body)
				mu.Lock()
				errorsList = append(errorsList, fmt.Errorf("导出任务ID %s 失败，响应: %s", taskID, string(body)))
				mu.Unlock()
				logger.Printf("导出任务ID %s 失败，响应: %s", taskID, string(body))
				return
			}

			// 获取文件名并进行URL解码
			contentDisposition := resp.Header.Get("Content-Disposition")
			filename := fmt.Sprintf("%s_report_%s.xlsx", taskName, taskID)
			if strings.Contains(contentDisposition, "filename=") {
				filename = strings.Split(contentDisposition, "filename=")[1]
				filename = strings.Trim(filename, "\"")
				decodedFilename, err := url.QueryUnescape(filename)
				if err == nil {
					filename = decodedFilename
				}
			}

			// 保存文件
			outFile, err := os.Create(filename)
			if err != nil {
				mu.Lock()
				errorsList = append(errorsList, fmt.Errorf("创建文件 %s 失败: %v", filename, err))
				mu.Unlock()
				logger.Printf("创建文件 %s 失败: %v", filename, err)
				return
			}
			defer outFile.Close()

			_, err = io.Copy(outFile, resp.Body)
			if err != nil {
				mu.Lock()
				errorsList = append(errorsList, fmt.Errorf("保存文件 %s 失败: %v", filename, err))
				mu.Unlock()
				logger.Printf("保存文件 %s 失败: %v", filename, err)
				return
			}

			color.Green("报告已成功导出为: %s", filename)
			logger.Printf("报告已成功导出为: %s", filename)
		}(taskID)
	}

	wg.Wait()

	if len(errorsList) > 0 {
		color.Red("部分报告导出失败:")
		for _, e := range errorsList {
			color.Red("  - %v", e)
			logger.Printf("报告导出失败: %v", e)
		}
	} else {
		color.Green("[+] 所有报告已成功导出")
		logger.Println("所有报告已成功导出")
	}

	return nil
}

// 分配任务到多服务器并控制并发速度
func distributeTasks(servers []ServerConfig, tokens []string, taskName string, targets []string) {
	serverCount := len(servers)
	targetCount := len(targets)

	if serverCount == 0 {
		color.Red("未配置任何服务器，无法分配任务。")
		logger.Println("未配置任何服务器，无法分配任务。")
		return
	}

	// 计算每台服务器应分配的任务数量
	base := targetCount / serverCount
	remainder := targetCount % serverCount

	// 创建任务分配
	taskAssignments := make([][]string, serverCount)
	current := 0
	for i := 0; i < serverCount; i++ {
		assignCount := base
		if i < remainder {
			assignCount++ // 将多余的任务分配给前几台服务器
		}
		if assignCount > 0 {
			taskAssignments[i] = targets[current : current+assignCount]
			current += assignCount
		}
	}

	var wg sync.WaitGroup
	var mu sync.Mutex
	var errorsList []error

	// 定义并发限制，例如每次最多10个并发任务
	const maxConcurrentTasks = 10
	semaphore := make(chan struct{}, maxConcurrentTasks)

	for i, server := range servers {
		assignedTasks := taskAssignments[i]
		if len(assignedTasks) == 0 {
			continue
		}
		wg.Add(1)
		go func(server ServerConfig, token string, tasks []string) {
			defer wg.Done()
			for _, target := range tasks {
				semaphore <- struct{}{} // 获取一个信号量
				wg.Add(1)
				go func(target string) {
					defer wg.Done()
					defer func() { <-semaphore }() // 释放信号量

					err := addTask(server.URL, token, taskName, target, &server)
					if err != nil {
						mu.Lock()
						errorsList = append(errorsList, fmt.Errorf("服务器 %s 目标 %s: %v", server.Name, target, err))
						mu.Unlock()
						color.Red("[-] 服务器 %s 任务 %s 创建失败: %v", server.Name, target, err)
						logger.Printf("服务器 %s 任务 %s 创建失败: %v", server.Name, target, err)
					} else {
						color.Green("[+] 服务器 %s 任务已成功创建：%s 目标：%s", server.Name, taskName, target)
						logger.Printf("服务器 %s 任务已成功创建：%s 目标：%s", server.Name, taskName, target)
					}
				}(target)
			}
		}(server, tokens[i], assignedTasks)
	}

	wg.Wait()

	if len(errorsList) > 0 {
		color.Red("部分任务创建失败:")
		for _, e := range errorsList {
			color.Red("  - %v", e)
			logger.Printf("任务创建失败: %v", e)
		}
	} else {
		color.Green("[+] 所有任务已成功创建")
		logger.Println("所有任务已成功创建")
	}
}

// 输出扫描配置（单服务器）
func displayScanConfig(config *ServerConfig) {
	color.Cyan("\n🔹 扫描配置:")
	fmt.Printf("  域名爆破: %t\n", config.DomainBrute)
	fmt.Printf("  域名查询插件: %t\n", config.DNSQueryPlugin)
	fmt.Printf("  端口扫描: %t\n", config.PortScan)
	fmt.Printf("  跳过CDN: %t\n", config.SkipScanCDNIP)
	fmt.Printf("  服务识别: %t\n", config.ServiceDetection)
	fmt.Printf("  操作系统识别: %t\n", config.OSDetection)
	fmt.Printf("  站点识别: %t\n", config.SiteIdentify)
	fmt.Printf("  文件泄露扫描: %t\n", config.FileLeak)
	fmt.Printf("  DNS字典智能生成: %t\n", config.AltDNS)
	fmt.Printf("  站点爬虫: %t\n", config.SiteSpider)
	fmt.Printf("  站点截图: %t\n", config.SiteCapture)
	fmt.Printf("  ARL 历史查询: %t\n", config.ARLSearch)
}

// 输出扫描配置（多服务器）
func displayScanConfigs(servers []ServerConfig) {
	color.Cyan("\n🔹 扫描配置:")
	for _, config := range servers {
		fmt.Printf("\n🔸 服务器: %s\n", config.Name)
		fmt.Printf("  域名爆破: %t\n", config.DomainBrute)
		fmt.Printf("  域名查询插件: %t\n", config.DNSQueryPlugin)
		fmt.Printf("  端口扫描: %t\n", config.PortScan)
		fmt.Printf("  跳过CDN: %t\n", config.SkipScanCDNIP)
		fmt.Printf("  服务识别: %t\n", config.ServiceDetection)
		fmt.Printf("  操作系统识别: %t\n", config.OSDetection)
		fmt.Printf("  站点识别: %t\n", config.SiteIdentify)
		fmt.Printf("  文件泄露扫描: %t\n", config.FileLeak)
		fmt.Printf("  DNS字典智能生成: %t\n", config.AltDNS)
		fmt.Printf("  站点爬虫: %t\n", config.SiteSpider)
		fmt.Printf("  站点截图: %t\n", config.SiteCapture)
		fmt.Printf("  ARL 历史查询: %t\n", config.ARLSearch)
	}
}

// 获取任务列表并进行美化格式输出
func getTaskList(url, token string, size int) {
	// 构建带有 size 参数的URL
	requestURL := fmt.Sprintf("%sapi/task/?size=%d&page=1", url, size)
	req, err := http.NewRequest("GET", requestURL, nil)
	if err != nil {
		color.Red("获取任务列表失败: %v", err)
		logger.Printf("获取任务列表失败: %v", err)
		return
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Token", token)

	resp, err := httpClient.Do(req)
	if err != nil {
		color.Red("请求失败: %v", err)
		logger.Printf("请求失败: %v", err)
		return
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		color.Red("获取任务列表失败: %s", string(body))
		logger.Printf("获取任务列表失败: %s", string(body))
		return
	}

	// 解析任务列表
	var taskList TaskListResponse
	err = json.Unmarshal(body, &taskList)
	if err != nil {
		color.Red("解析任务列表失败: %v", err)
		logger.Printf("解析任务列表失败: %v", err)
		return
	}

	if len(taskList.Items) == 0 {
		color.Yellow("当前没有任务。")
		return
	}

	// 使用 tablewriter 设置更美观的表格格式化输出
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"任务名", "站点统计", "域名统计", "状态", "开始时间", "结束时间", "任务ID"})
	table.SetBorder(true)
	table.SetRowLine(true) // 使用实线
	table.SetHeaderColor(
		tablewriter.Colors{tablewriter.FgGreenColor},
		tablewriter.Colors{tablewriter.FgGreenColor},
		tablewriter.Colors{tablewriter.FgGreenColor},
		tablewriter.Colors{tablewriter.FgGreenColor},
		tablewriter.Colors{tablewriter.FgGreenColor},
		tablewriter.Colors{tablewriter.FgGreenColor},
		tablewriter.Colors{tablewriter.FgGreenColor},
	)
	table.SetColumnColor(
		tablewriter.Colors{tablewriter.FgCyanColor},
		tablewriter.Colors{tablewriter.FgHiGreenColor},
		tablewriter.Colors{tablewriter.FgHiGreenColor},
		tablewriter.Colors{tablewriter.FgHiYellowColor},
		tablewriter.Colors{tablewriter.FgHiMagentaColor},
		tablewriter.Colors{tablewriter.FgHiMagentaColor},
		tablewriter.Colors{tablewriter.FgHiRedColor},
	)
	table.SetHeaderAlignment(tablewriter.ALIGN_CENTER)
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetAutoWrapText(false)

	// 遍历任务列表并添加到表格
	for _, task := range taskList.Items {
		statisticsSites := fmt.Sprintf("%d", task.Statistic.SiteCnt)
		statisticsDomains := fmt.Sprintf("%d", task.Statistic.DomainCnt)
		startTime := formatTime(task.StartTime)
		endTime := formatTime(task.EndTime)
		status := getStatus(task.Status)

		row := []string{
			task.Name,
			statisticsSites,
			statisticsDomains,
			status,
			startTime,
			endTime,
			task.ID,
		}
		table.Append(row)
	}

	// 渲染表格
	table.Render()
}

// 格式化时间，保持简洁
func formatTime(timeStr string) string {
	if timeStr == "-" || timeStr == "" {
		return "-"
	}

	parsedTime, err := time.Parse("2006-01-02 15:04:05", timeStr)
	if err != nil {
		return timeStr // 返回原始时间格式
	}

	return parsedTime.Format("2006/01/02 15:04")
}

// 获取任务状态并返回对应的颜色
func getStatus(status string) string {
	switch status {
	case "done":
		return color.GreenString("已完成")
	case "stop":
		return color.RedString("已停止")
	case "waiting":
		return color.YellowString("等待中")
	case "running":
		return color.BlueString("正在运行")
	default:
		return color.WhiteString(status)
	}
}

// 验证配置文件中的必填字段
func validateConfig(config *Config) error {
	for _, server := range config.Servers {
		if server.Name == "" {
			return fmt.Errorf("某个服务器的别名未配置")
		}
		if server.URL == "" {
			return fmt.Errorf("服务器 %s 的URL未配置", server.Name)
		}
		if server.Username == "" {
			return fmt.Errorf("服务器 %s 的用户名未配置", server.Name)
		}
		if server.Password == "" {
			return fmt.Errorf("服务器 %s 的密码未配置", server.Name)
		}
	}
	return nil
}

// 打印成功消息并记录日志
func printSuccessMessage(taskName, target, serverName string) {
	color.Green("[+] 任务已成功创建：%s 目标：%s (服务器: %s)", taskName, target, serverName)
	logger.Printf("任务已成功创建：%s 目标：%s (服务器: %s)", taskName, target, serverName)
}

// 打印错误消息并记录日志
func printErrorMessage(err error, serverName string) {
	color.Red("[-] %s 错误: %v", serverName, err)
	logger.Printf("服务器 %s 错误: %v", serverName, err)
}
